import cv2
import numpy as np
import os
from datetime import datetime
import matplotlib.pyplot as plt

def preprocess_grid(image_path: str, output_dir="static/preprocess") -> str:
    os.makedirs(output_dir, exist_ok=True)

    img = cv2.imread(image_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)

    # 保存二值化图像
    binary_filename = f"binary_{datetime.now().strftime('%Y%m%d%H%M%S')}.png"
    binary_path = os.path.join(output_dir, "binary")
    binary_path = os.path.join(binary_path, binary_filename)
    cv2.imwrite(binary_path, binary)
    print(f"[DEBUG] 二值化图像已保存：{binary_path}")

    height, width = binary.shape

    # === 横向投影：获取每一行中线 ===
    horizontal_sum = np.sum(binary, axis=1)
    row_positions = []
    threshold = np.max(horizontal_sum) * 0.22

    in_text = False
    for i, val in enumerate(horizontal_sum):
        if val > threshold and not in_text:
            start = i
            in_text = True
        elif val <= threshold and in_text:
            end = i
            in_text = False
            if end - start > 5:
                row_positions.append((start, end))

    row_centers = [(start + end) // 2 for start, end in row_positions]
    row_lines = set()
    for i in range(len(row_centers) - 1):
        mid = (row_centers[i] + row_centers[i + 1]) // 2
        row_lines.add(mid)

    # 可选添加首尾
    if len(row_centers) >= 2:
        d = row_centers[1] - row_centers[0]
        row_lines.add(row_centers[0] - d // 2)
        row_lines.add(row_centers[-1] + d // 2)

    # === 垂直投影：找最宽的两个 valley 的中线 ===
    vertical_sum = np.sum(binary, axis=0)
    smooth = cv2.GaussianBlur(vertical_sum.astype(np.float32).reshape(1, -1), (1, 21), 0).flatten()

    # 找连续低于阈值的区间
    threshold = np.max(smooth) * 0.1
    valleys = []
    in_valley = False
    start = 0

    for i, val in enumerate(smooth):
        if val < threshold and not in_valley:
            in_valley = True
            start = i
        elif val >= threshold and in_valley:
            end = i
            in_valley = False
            width_valley = end - start
            if width_valley > 10:
                valleys.append((start, end, width_valley))

    # 取最宽的两个 valley 中心
    valleys = sorted(valleys, key=lambda x: x[2], reverse=True)[:2]
    valley_mids = sorted([(start + end) // 2 for start, end, _ in valleys])

    # 如果不够两个，fallback
    if len(valley_mids) < 2:
        valley_mids = [int(width * 0.4), int(width * 0.75)]

    col_lines = [0] + valley_mids + [width - 1]

    # === 绘图 ===
    output_img = img.copy()
    for y in sorted(row_lines):
        cv2.line(output_img, (0, y), (width, y), (0, 0, 0), 2)
    for x in col_lines:
        cv2.line(output_img, (x, 0), (x, height), (0, 0, 0), 2)

    # === 保存主图 ===
    filename = f"grid_3col_{datetime.now().strftime('%Y%m%d%H%M%S')}.png"
    grid_path = os.path.join(output_dir, "grid")
    output_path = os.path.join(grid_path, filename)
    cv2.imwrite(output_path, output_img)
    print(f"[DEBUG] 三列辅助线已保存：{output_path}")

    # # === 保存垂直投影调试图 ===
    debug_plot_path = os.path.join(output_dir, "vertical_projection")
    debug_plot_path = os.path.join(debug_plot_path, f"vertical_projection_{datetime.now().strftime('%Y%m%d%H%M%S')}.png")
    plt.figure(figsize=(12, 4))
    plt.plot(smooth, label='Smoothed Vertical Projection')
    for x in valley_mids:
        plt.axvline(x=x, color='red', linestyle='--', label='Valley Mid' if x == valley_mids[0] else None)
    plt.title("Vertical Projection with Valley Centers")
    plt.xlabel("X (columns)")
    plt.ylabel("Pixel sum")
    plt.legend()
    plt.tight_layout()
    plt.savefig(debug_plot_path)
    plt.close()
    print(f"[DEBUG] 垂直投影调试图已保存：{debug_plot_path}")

    return output_path
